package com.ds.infra.mq.core.api;

import com.ds.infra.mq.core.api.exception.ConsumeException;
import com.ds.infra.mq.core.constants.MQConstant;
import com.ds.infra.mq.core.protocol.Message;
import com.ds.infra.mq.core.protocol.MessageHeader;
import com.ds.infra.mq.core.protocol.MessageType;
import com.ds.infra.mq.core.protocol.ProtocolCodec;
import io.netty.bootstrap.Bootstrap;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import lombok.extern.slf4j.Slf4j;

import java.nio.charset.StandardCharsets;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 作用：MQConsumer类实现了Consumer接口，用于从Broker中消费消息。
 *
 * @author WeiShaoying
 */
@Slf4j
public class MQConsumer implements Consumer {

    private final String brokerHost;
    private final int brokerPort;
    private final String clientId;
    private final AtomicLong requestIdGenerator = new AtomicLong(0);
    private final BlockingQueue<byte[]> messageQueue = new LinkedBlockingQueue<>(100);
    private final long pollTimeoutMs = 5000;

    private EventLoopGroup group;
    private Channel channel;

    public MQConsumer(String brokerHost, int brokerPort, String clientId) {
        this.brokerHost = brokerHost;
        this.brokerPort = brokerPort;
        this.clientId = clientId != null ? clientId : "consumer-" + System.currentTimeMillis();
    }

    public void connect() throws InterruptedException {
        if (channel != null && channel.isActive()) {
            // 如果已经连接，则直接返回
            return;
        }
        group = new NioEventLoopGroup(1);
        Bootstrap bootstrap = new Bootstrap();
        bootstrap.group(group)
                .channel(NioSocketChannel.class)
                .handler(new ChannelInitializer<SocketChannel>() {
                    @Override
                    protected void initChannel(SocketChannel ch) {
                        ch.pipeline()
                                .addLast(new ProtocolCodec.Decoder())
                                .addLast(new ProtocolCodec.Encoder())
                                .addLast(new ConsumerHandler(messageQueue));
                    }
                });
        ChannelFuture future = bootstrap.connect(brokerHost, brokerPort).sync();
        channel = future.channel();
        log.info("Consumer [{}] connected to {}:{}", clientId, brokerHost, brokerPort);
    }

    @Override
    public byte[] poll(String topic) throws ConsumeException {
        return poll(topic, 0); // 默认分区0
    }

    @Override
    public byte[] poll(String topic, int partition) throws ConsumeException {
        try {
            MessageHeader header = new MessageHeader();
            header.setMessageType(MessageType.CONSUME_REQUEST.getCode());
            header.setRequestId(requestIdGenerator.incrementAndGet());
            header.setClientId(clientId);

            Message msg = new Message();
            msg.setHeader(header);
            msg.setBody((topic + ":" + partition).getBytes(StandardCharsets.UTF_8));

            ChannelFuture future = channel.writeAndFlush(msg).sync();
            if (!future.isSuccess()) {
                throw new ConsumeException("Consume send poll request failed");
            }
            // 阻塞等待消息，超时则返回null
            byte[] messageBytes = messageQueue.poll(pollTimeoutMs, TimeUnit.MILLISECONDS);

            // 安全处理 null 和 EMPTY 响应
            if (messageBytes == null || messageBytes.length == 0 || MQConstant.EMPTY_STR.equals(new String(messageBytes))) {
                return null;
            }
            return messageBytes;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new ConsumeException("Poll interrupted", e);
        } catch (Exception e) {
            throw new ConsumeException("Poll failed", e);
        }
    }

    @Override
    public void subscribe(String topic) {
        // 实际订阅逻辑在poll中实现
        log.info("Consumer [{}] subscribed to topic: {}", clientId, topic);
    }

    @Override
    public void unsubscribe(String topic) {
        // 可添加取消订阅逻辑
        log.info("Consumer [{}] unsubscribed from topic: {}", clientId, topic);
    }

    @Override
    public void close() {
        if (channel != null) {
            channel.close().awaitUninterruptibly();
        }
        if (group != null) {
            group.shutdownGracefully();
        }
        log.info("Consumer [{}] closed", clientId);
    }

    @ChannelHandler.Sharable
    private static class ConsumerHandler extends SimpleChannelInboundHandler<Message> {

        private final BlockingQueue<byte[]> messageQueue;

        public ConsumerHandler(BlockingQueue<byte[]> messageQueue) {
            this.messageQueue = messageQueue;
        }

        @Override
        protected void channelRead0(ChannelHandlerContext ctx, Message msg) {
            if (msg.getHeader().getMessageType() == MessageType.CONSUME_RESPONSE.getCode()) {
                // 消费响应，将消息放入队列中
                messageQueue.offer(msg.getBody());
            }
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
            log.error("Consumer handler error", cause);
            ctx.close();
        }
    }
}
